import numpy as np
import h5py
from scipy.optimize import minimize
from tools.Cal_Z import Cal_Z
from tools.read_hdf5_find_powers import read_hdf5_find_powers
from tools.mag_phase_to_IQ import mag_phase_to_IQ
from tools.calculate_deri import calculate_deri
from tools.Fdecay_fhalf import Fdecay_fhalf

"""
此程序用于计算腔体校准参数a，并保存相关电压数据。
主要功能：
1. 从HDF5文件读取原始电压数据
2. 计算最优的校准参数a
3. 将计算结果和电压数据保存到新的HDF5文件中
"""

# 系统参数定义
RQ = 8893      # 电阻-品质因数乘积，单位：Ω
Ql = 76.7      # 外部品质因数，无量纲
beta = 1.2     # 耦合系数，无量纲
Fs = 117.402e6 / 14  # 采样率，单位：Hz

def calculate_y_a(a, Z, Vc, Vf, Vr):
    """
    计算y_a参数，用于优化过程中的计算。
    
    参数:
        a (list): 包含两个元素的列表 [real, imag]，表示复数a的实部和虚部
        Z (complex): 腔体阻抗
        Vc (ndarray): 腔体电压
        Vf (ndarray): 前向电压
        Vr (ndarray): 反射电压
    
    返回:
        ndarray: 计算得到的y_a值
    """
    w_half = Fdecay_fhalf(Vc, Vf, Fs)  # 计算半高宽
    a_complex = complex(a[0], a[1])     # 构造复数形式的a
    Vf_cal = a_complex * Vf + a_complex / Z * Vr  # 计算校准后的前向电压
    rho = 2 * np.abs((beta / (1 + beta)) * Vf_cal)  # 计算ρ值
    return w_half * (rho * np.cos(np.angle(Vc) - np.angle(Vf_cal)) - np.abs(Vc))

def error_function(a, Z, Vc, Vf, Vr):
    """
    计算误差函数，用于优化过程。
    
    参数:
        a (list): 待优化的参数 [real, imag]
        Z (complex): 腔体阻抗
        Vc, Vf, Vr (ndarray): 腔体、前向和反射电压
    
    返回:
        float: 误差函数值
    """
    dVc_dt = calculate_deri(np.abs(Vc), Fs)  # 计算腔体电压的导数
    y_a = calculate_y_a(a, Z, Vc, Vf, Vr)    # 计算y_a
    return np.sum((dVc_dt - y_a) ** 2)       # 返回误差平方和

def save_to_hdf5(filename, Vc, Vf, Vr, a_opt):
    """
    将计算结果保存到HDF5文件。
    
    参数:
        filename (str): 输出文件名
        Vc, Vf, Vr (ndarray): 腔体、前向和反射电压数据
        a_opt (complex): 优化后的a参数
    """
    with h5py.File(filename, 'w') as f:
        # 创建电压数据组
        voltage_group = f.create_group('voltages')
        
        # 分别保存电压的实部和虚部
        voltage_group.create_dataset('Vc_real', data=Vc.real)
        voltage_group.create_dataset('Vc_imag', data=Vc.imag)
        voltage_group.create_dataset('Vf_real', data=Vf.real)
        voltage_group.create_dataset('Vf_imag', data=Vf.imag)
        voltage_group.create_dataset('Vr_real', data=Vr.real)
        voltage_group.create_dataset('Vr_imag', data=Vr.imag)
        
        # 保存优化后的a参数
        param_group = f.create_group('parameters')
        param_group.create_dataset('a_real', data=a_opt.real)
        param_group.create_dataset('a_imag', data=a_opt.imag)

def main():
    """
    主函数，执行完整的计算和保存流程。
    """
    # 加载原始数据
    filename = r'D:\Beam phase calibration\MEBT_50us_scan_beam_current_3mA_63mA_B1_B2\AFF\63mA_70mm_aff\RFQ-Bunchers_3-60mA_50us_1628_2023May25.hdf5'
    raw_powers = read_hdf5_find_powers(filename, 1)  # 读取原始功率数据
    voltages_IQ = mag_phase_to_IQ(raw_powers, RQ, Ql, beta)  # 转换为IQ形式

    # 将IQ分量转换为复数形式的电压数据
    Vc = voltages_IQ['Vcav_I'] + 1j * voltages_IQ['Vcav_Q']  # 腔体电压
    Vf = voltages_IQ['Vfor_I'] + 1j * voltages_IQ['Vfor_Q']  # 前向电压
    Vr = voltages_IQ['Vrefl_I'] + 1j * voltages_IQ['Vrefl_Q']  # 反射电压

    # 计算阻抗和优化a参数
    Z = Cal_Z(Vf, Vr)  # 计算腔体阻抗
    result = minimize(error_function, [0, 0.0], args=(Z, Vc, Vf, Vr), method='Nelder-Mead')
    a_opt = complex(result.x[0], result.x[1])  # 获取优化后的a参数

    # 打印计算结果
    print(f"\n计算结果:")
    print(f"a = {a_opt:.6f}")
    print(f"Z = {Z:.6f}")
    print(f"最小误差 = {result.fun:.6f}")

    # 保存数据到HDF5文件
    output_filename = 'voltage_data_with_a.h5'
    save_to_hdf5(output_filename, Vc, Vf, Vr, a_opt)
    print(f"\n数据已保存到文件: {output_filename}")

if __name__ == '__main__':
    main() 